
#include <ff_wrapper_vk.h>
#include <system_log.h>
#include <cstring>
#include <unordered_map>

namespace FFSdkUnity::Vulkan {

static PFN_vkGetInstanceProcAddr vkGetInstanceProcAddr = nullptr;
static PFN_vkCreateImageView vkCreateImageView = nullptr;
static PFN_vkDestroyImageView vkDestroyImageView = nullptr;

static bool LoadVulkanFunc(VkInstance instance, PFN_vkGetInstanceProcAddr getInstanceProcAddr)
{
    if (instance == VK_NULL_HANDLE || getInstanceProcAddr == nullptr) {
        return false;
    }
    vkGetInstanceProcAddr = getInstanceProcAddr;
    vkCreateImageView = (PFN_vkCreateImageView)vkGetInstanceProcAddr(instance, "vkCreateImageView");
    vkDestroyImageView = (PFN_vkDestroyImageView)vkGetInstanceProcAddr(instance, "vkDestroyImageView");
    return true;
}

static void UnloadVulkanFunc()
{
    vkGetInstanceProcAddr = nullptr;
    vkCreateImageView = nullptr;
    vkDestroyImageView = nullptr;
}

constexpr static uint8_t FRAMES_IN_FLIGHT = 2;

std::unordered_map<FFVKImage, VkImageView> registered_images;

FFVKImage g_unity_color = FF_NULL_HANDLE;
FFVKImage g_unity_zs = FF_NULL_HANDLE;
FFVKImage g_unity_predict = FF_NULL_HANDLE;

static std::unordered_map<FFUnityZSFormat, VkFormat> g_depth_remap = {
    {D32_SFLOAT_S8_UINT, VK_FORMAT_D32_SFLOAT_S8_UINT},
    {D24_UNORM_S8_UINT, VK_FORMAT_D24_UNORM_S8_UINT},
    {D16_UNORM_S8_UINT, VK_FORMAT_D16_UNORM_S8_UINT}
};

static std::unordered_map<FFUnityColorFormat, VkFormat> g_color_remap = {
    {R8G8B8A8_UNORM, VK_FORMAT_R8G8B8A8_UNORM}
};

bool get_depth_fmt(FFUnityZSFormat fmt_unity, VkFormat& fmt_vk)
{
    auto it = g_depth_remap.find(fmt_unity);
    if (it == g_depth_remap.end()) {
        return false;
    }
    fmt_vk = it->second;
    return true;
}

bool get_color_fmt(FFUnityColorFormat fmt_unity, VkFormat& fmt_vk)
{
    auto it = g_color_remap.find(fmt_unity);
    if (it == g_color_remap.end()) {
        return false;
    }
    fmt_vk = it->second;
    return true;
}

bool FFWrapperVK::Init(IUnityGraphicsVulkan* unity_graphics) noexcept
{
    FFLOGI("Sdk::Unity::FFWrapperVK::Init - Start");
    if (ffInstance_ != FF_NULL_HANDLE) {
        FFLOGI("Sdk::Unity::FFWrapperVK::Init - Success, as already init");
        return true;
    }

    if (!unity_graphics) {
        FFLOGE("Sdk::Unity::FFWrapperVK::Init - IUnityGraphicsVulkan is null");
        return false;
    }

    // init Vulkan API functions
    unity_graphics_ = unity_graphics;
    auto vk_instance = unity_graphics->Instance();
    if (!vk_instance.getInstanceProcAddr) {
        FFLOGE("Sdk::Unity::FFWrapperVK::Init - getInstanceProcAddr is null");
        return false;
    }

    bool ret = LoadVulkanFunc(vk_instance.instance, vk_instance.getInstanceProcAddr);
    if (!ret) {
        FFLOGE("Sdk::Unity::FFWrapperVK::Init - LoadVulkanFunc failed");
        return false;
    }

    if (!LoadFrameFlow()) {
        FFLOGE("Sdk::Unity::FFWrapperVK::Init - Can't LoadFrameFlow");
        if (!Destroy()) {
            FFLOGE("Sdk::Unity::FFWrapperVK::Init - Can't destroy wrapper resources (branch 2)");
        }
        return false;
    }
    FFLOGI("Sdk::Unity::FFWrapperVK::Init - Load frameflow so success");

    // Load SDK functions
#define FF_FUNCTION(func)                                                                                              \
    if ((func) = reinterpret_cast<decltype(func)>(frameflowLib_.GetFunction(#func)); !(func)) {                        \
        FFLOGE("Sdk::Unity::FFWrapperVK::FindSDKFunctionsLib: function %s not found", #func);                          \
        return false;                                                                                                  \
    }
#include <ff_function/ff_vk_functions.in>
#undef FF_FUNCTION
    FFLOGI("Sdk::Unity::FFWrapperVK::Init - Load frameflow functions success");

    UnityVulkanPluginEventConfig cfg{};
    cfg.graphicsQueueAccess = kUnityVulkanGraphicsQueueAccess_DontCare;
    cfg.renderPassPrecondition = kUnityVulkanRenderPass_EnsureOutside;
    cfg.flags = kUnityVulkanEventConfigFlag_ModifiesCommandBuffersState;
                //| kUnityVulkanEventConfigFlag_EnsurePreviousFrameSubmission;
    unity_graphics_->ConfigureEvent(1, &cfg);

    FFVKInstanceCreateInfo ffCreateInfo{};
    ffCreateInfo.sType                   = FF_STRUCTURE_TYPE_INSTANCE_CREATE_INFO;
    ffCreateInfo.pNext                   = nullptr;
    ffCreateInfo.majorVersion            = FFVK_API_VERSION_MAJOR;
    ffCreateInfo.minorVersion            = FFVK_API_VERSION_MINOR;
    ffCreateInfo.vkInstance              = vk_instance.instance;
    ffCreateInfo.vkPhysicalDevice        = vk_instance.physicalDevice;
    ffCreateInfo.vkDevice                = vk_instance.device;
    ffCreateInfo.framesInFlight          = static_cast<uint8_t>(FRAMES_IN_FLIGHT);
    ffCreateInfo.pfnVulkanLoaderFunction = vk_instance.getInstanceProcAddr;
    ffCreateInfo.pfnErrorReportCallback  = &FFWrapperVK::SDKErrorCallback;
    ffCreateInfo.pUserData               = this;

    ffInstance_ = FrameFlowVK_Create(&ffCreateInfo);
    ffInstanceIsExternal_ = false;

    if (ffInstance_ == FF_NULL_HANDLE) {
        FFLOGE("Sdk::Unity::FFWrapperVK::Init - Can't create FrameFlow instance");
        if (!Destroy()) {
            FFLOGE("Sdk::Unity::FFWrapperVK::Init - Can't destroy wrapper resources (branch 4)");
        }
        return false;
    }
    FFLOGI("Sdk::Unity::FFWrapperVK::Init - Success 1");
    return true;
}

bool FFWrapperVK::Destroy() noexcept
{
    FFLOGI("Sdk::Unity::FFWrapperVK::Destroy - Start");
    UnRegisterImage(g_unity_color);
    UnRegisterImage(g_unity_zs);
    UnRegisterImage(g_unity_predict);

    if (ffInstance_ != FF_NULL_HANDLE && !ffInstanceIsExternal_) {
        FrameFlowVK_Destroy(ffInstance_);
    }
    ffInstance_ = FF_NULL_HANDLE;

    UnloadVulkanFunc();
    unity_graphics_ = nullptr;
    return frameflowLib_.Unload();
}

FFResult FFWrapperVK::OnActivate() const noexcept
{
    FFLOGI("Sdk::Unity::FFWrapperVK::OnActivate - Start");
    return IsInit() ? FrameFlowVK_Activate(ffInstance_) : FFResult::FF_ERROR_NOT_INITIALIZED;
}

FFResult FFWrapperVK::OnDeactivate() const noexcept
{
    return IsInit() ? FrameFlowVK_Deactivate(ffInstance_) : FFResult::FF_ERROR_NOT_INITIALIZED;
}

FFBool FFWrapperVK::OnIsActive() const noexcept
{
    return IsInit() ? FrameFlowVK_IsActive(ffInstance_) : FF_FALSE;
}

FFResult FFWrapperVK::OnSetPredictionAlgorithm(FFPredictionAlgorithm algorithm, bool extrapolation) const noexcept
{
    FFLOGI("Sdk::Unity::FFWrapperVK::OnSetPredictionAlgorithm - Start");
    if (!IsInit()) {
        return FFResult::FF_ERROR_NOT_INITIALIZED;
    }

    FFPredictionAlgorithmInfo info{};
    info.sType = FF_STRUCTURE_TYPE_PREDICTION_ALGORITHM_INFO;
    info.pNext = nullptr;
    info.algorithm = algorithm;
    info.extrapolation = extrapolation;

    auto ret = FrameFlowVK_SetPredictionAlgorithm(ffInstance_, &info, nullptr);
    FFLOGI("Sdk::Unity::FFWrapperVK::OnSetPredictionAlgorithm - end, return %x", ret);
    return ret;
}

FFResult FFWrapperVK::OnSetStencilRemapInfo(FFUnityRemapTable const *unity_table) noexcept
{
    FFLOGI("Sdk::Unity::FFWrapperVK::OnSetStencilRemapInfo - Start");
    if (!IsInit()) {
        return FFResult::FF_ERROR_NOT_INITIALIZED;
    }

    if (!unity_table) {
        FFLOGE("FFWrapperVK::OnSetStencilRemapInfo. Unity data is null.");
        return FFResult::FF_ERROR_INSUFFICIENT_INPUT_DATA;
    }    

    auto pRemapTable = unity_table->remapTable;

    FFStencilRemapInfo info{};
    info.sType = FF_STRUCTURE_TYPE_STENCIL_REMAP_INFO;
    info.pNext = nullptr;
    auto const memcpyResult = std::memcpy(info.remapTable, pRemapTable, sizeof(info.remapTable));
    if (memcpyResult == nullptr) {
        return FFResult::FF_ERROR_INTERNAL_ERROR;
    }

    auto ret = FrameFlowVK_SetStencilRemapInfo(ffInstance_, &info);
    FFLOGI("Sdk::Unity::FFWrapperVK::OnSetStencilRemapInfo - end, return %x", ret);
    return ret;
}

bool FFWrapperVK::RegisterImage(void* unity_handle, bool zs_image, const FFVKImageSync& initial_sync, FFVKImage& ff_image)
{
    FFLOGI("Sdk::Unity::FFWrapperVK::RegisterImage - Start");
    if (unity_handle == nullptr) {
        FFLOGE("FFWrapperVK::RegisterImage. Unity texture is null.");
        return false;
    }

    if (ff_image != FF_NULL_HANDLE)
        return true;

    UnityVulkanImage unity_color{};
    bool result = unity_graphics_->AccessTexture(
           unity_handle,
           UnityVulkanWholeImage,
           initial_sync.layout,
           initial_sync.stages,
           initial_sync.accessMask,
           kUnityVulkanResourceAccess_PipelineBarrier, &unity_color);

    if (!result) {
        FFLOGE("FFWrapperVK::RegisterImage. Cant access unity texture.");
        return false;
    }

    VkImageAspectFlags image_aspects = 0;
    if (zs_image) {
        image_aspects = VK_IMAGE_ASPECT_DEPTH_BIT;
        image_aspects |= VK_IMAGE_ASPECT_STENCIL_BIT;        
    } else {
        image_aspects = VK_IMAGE_ASPECT_COLOR_BIT;
    }

    VkImageViewCreateInfo view_info {};
    view_info.sType = VK_STRUCTURE_TYPE_IMAGE_VIEW_CREATE_INFO;
    view_info.viewType = VK_IMAGE_VIEW_TYPE_2D;
    view_info.format = unity_color.format;
    view_info.subresourceRange = { image_aspects, 0, 1, 0, 1 };
    view_info.image = unity_color.image;

    auto vk_instance = unity_graphics_->Instance();
    VkImageView image_view = VK_NULL_HANDLE;
    auto create_result = vkCreateImageView(vk_instance.device, &view_info, nullptr, &image_view);
    if (create_result != VK_SUCCESS) {
        FFLOGE("FFWrapperVK::RegisterImage. Cant create view for unity texture");
        return false;
    }

    ff_image = FrameFlowVK_RegisterImage(ffInstance_, unity_color.image, image_view);
    if (ff_image != FF_NULL_HANDLE) {
        registered_images[ff_image] = image_view;
    } else {
        vkDestroyImageView(vk_instance.device, image_view, nullptr);
    }

    FFLOGI("Sdk::Unity::FFWrapperVK::RegisterImage - end");
    return (ff_image != FF_NULL_HANDLE);
}

void FFWrapperVK::UnRegisterImage(FFVKImage& ff_image)
{
    FFLOGI("Sdk::Unity::FFWrapperVK::UnRegisterImage - Start");
    VkDevice device = VK_NULL_HANDLE;
    if (unity_graphics_)
        device = unity_graphics_->Instance().device;

    auto it = registered_images.find(ff_image);
    if (it != registered_images.end()) {
        if (ff_image != FF_NULL_HANDLE)
            FrameFlowVK_UnregisterImage(ffInstance_, ff_image);
        auto &view = it->second;
        if (view != VK_NULL_HANDLE) {
            vkDestroyImageView(device, view, nullptr);
            registered_images.erase(ff_image);
        }
    }
    ff_image = FF_NULL_HANDLE;
    FFLOGI("Sdk::Unity::FFWrapperVK::UnRegisterImage - end");
}

FFResult FFWrapperVK::OnSetInputDescription(FFUnityInputDescription const *unity_data) noexcept
{
    FFLOGI("Sdk::Unity::FFWrapperVK::OnSetInputDescription - Start");
    if (!IsInit())
        return FFResult::FF_ERROR_NOT_INITIALIZED;

    if (!unity_data) {
        FFLOGE("FFWrapperVK::OnSetInputDescription. Unity data is null.");
        return FFResult::FF_ERROR_INSUFFICIENT_INPUT_DATA;
    }

    FFLOGI("FFWrapperVK::OnSetInputDescription. init: w = %d, h = %d", unity_data->width, unity_data->height);

    VkFormat zs_fmt; // VK_FORMAT_D24_UNORM_S8_UINT
    if (!get_depth_fmt(unity_data->zs_fmt, zs_fmt)) {
        FFLOGE("FFWrapperVK::OnSetInputDescription. Unsupported zs format %d.", unity_data->zs_fmt);
        return FFResult::FF_ERROR_INCOMPATIBLE_INPUT_DATA;
    }
    
    VkFormat scene_fmt; // VK_FORMAT_R8G8B8A8_UNORM
    if (!get_color_fmt(unity_data->scene_fmt, scene_fmt)) {
        FFLOGE("FFWrapperVK::OnSetInputDescription. Unsupported scene color format %d.", unity_data->scene_fmt);
        return FFResult::FF_ERROR_INCOMPATIBLE_INPUT_DATA;
    }

    VkFormat output_fmt; // VK_FORMAT_R8G8B8A8_UNORM;
    if (!get_color_fmt(unity_data->output_fmt, output_fmt)) {
        FFLOGE("FFWrapperVK::OnSetInputDescription. Unsupported output color format %d.", unity_data->output_fmt);
        return FFResult::FF_ERROR_INCOMPATIBLE_INPUT_DATA;
    }

    FFCvvZSemanticInfo cvvZSemantic = {
        .sType = FF_STRUCTURE_TYPE_CVV_Z_SEMANTIC_INFO,
        .pNext = nullptr,
        .semantic = FF_CVV_Z_SEMANTIC_ZERO_TO_ONE_REVERSE_Z  // important as vulkan, Unity
    };

    FFVKOutputFormatInfo ffOutputFormat {
        .sType = FF_STRUCTURE_TYPE_SCENE_COLOR_FORMAT_INFO,
        .pNext = &cvvZSemantic,
        .format = output_fmt
    };

    FFVKSceneColorFormatInfo ffSceneColorFormat {
        .sType = FF_STRUCTURE_TYPE_SCENE_COLOR_FORMAT_INFO,
        .pNext = &ffOutputFormat,
        .format = scene_fmt
    };

    const FFVKDepthStencilFormatInfo ffZsformat {
        .sType = FF_STRUCTURE_TYPE_DEPTH_STENCIL_FORMAT_INFO,
        .pNext = &ffSceneColorFormat,
        .format = zs_fmt
    };

    const FFResolutionInfo2 resolutionInfo {
        .sType = FF_STRUCTURE_TYPE_RESOLUTION_INFO_2,
        .pNext = &ffZsformat,

        .sceneColorResolution = {
            .width = static_cast<uint32_t>(unity_data->width),
            .height = static_cast<uint32_t>(unity_data->height)
        }
    };

    const FFInputDescription ffInputDescription {
        .sType = FF_STRUCTURE_TYPE_INPUT_DESCRIPTION,
        .pNext = &resolutionInfo,
        .flags = 0U
    };

    return FrameFlowVK_SetInputDescription(ffInstance_, &ffInputDescription);
}

FFResult FFWrapperVK::OnProvidePerFrameData(FFUnityFrameData const *unity_data) noexcept
{
    FFLOGI("Sdk::Unity::FFWrapperVK::OnProvidePerFrameData - Start");
    unity_graphics_->EnsureOutsideRenderPass();

    FFVKImageSync initial_sync_color{};
    initial_sync_color.accessMask = VK_ACCESS_SHADER_READ_BIT;
    initial_sync_color.layout = VK_IMAGE_LAYOUT_COLOR_ATTACHMENT_OPTIMAL;
    initial_sync_color.stages = VK_PIPELINE_STAGE_FRAGMENT_SHADER_BIT;

    bool register_result = RegisterImage(unity_data->sceneColor, false, initial_sync_color, g_unity_color);
    if (!register_result) {
        FFLOGE("FFWrapperVK::OnProvidePerFrameData. Cant register input color texture.");
        return FF_ERROR_INTERNAL_ERROR;
    }

    FFVKImageSync initial_sync_zs{};
    initial_sync_zs.accessMask = VK_ACCESS_SHADER_READ_BIT;
    initial_sync_zs.layout = VK_IMAGE_LAYOUT_DEPTH_STENCIL_ATTACHMENT_OPTIMAL;
    initial_sync_zs.stages = VK_PIPELINE_STAGE_FRAGMENT_SHADER_BIT;

    register_result = RegisterImage(unity_data->depthStencil, true, initial_sync_zs, g_unity_zs);
    if (!register_result) {
        FFLOGE("FFWrapperVK::OnProvidePerFrameData. Cant register input depth texture.");
        return FF_ERROR_INTERNAL_ERROR;
    }

    FFVKPerFrameData my_data{};
    my_data.sType = FF_STRUCTURE_TYPE_PER_FRAME_DATA;
    my_data.pNext = nullptr;
    my_data.flags = unity_data->flags;
    for (int i = 0; i < 16; ++i) { // use memcpy
        my_data.viewProj.data[i] = unity_data->viewProj.data[i];
        my_data.invViewProj.data[i] = unity_data->invViewProj.data[i];
    }

    FFVKImageSync final_sync_color{};
    final_sync_color.accessMask = VK_ACCESS_SHADER_READ_BIT;
    final_sync_color.layout = VK_IMAGE_LAYOUT_TRANSFER_DST_OPTIMAL;
    final_sync_color.stages = VK_PIPELINE_STAGE_FRAGMENT_SHADER_BIT;    

    my_data.sceneColorImage = g_unity_color;
    my_data.sceneColorInitialSync = initial_sync_color;
    my_data.sceneColorFinalSync = final_sync_color;

    FFVKImageSync final_sync_zs{};
    final_sync_zs.accessMask = VK_ACCESS_SHADER_READ_BIT;
    final_sync_zs.layout = VK_IMAGE_LAYOUT_DEPTH_STENCIL_ATTACHMENT_OPTIMAL;
    final_sync_zs.stages = VK_PIPELINE_STAGE_FRAGMENT_SHADER_BIT;    

    my_data.depthStencilImage = g_unity_zs;
    my_data.depthStencilInitialSync = initial_sync_zs;
    my_data.depthStencilFinalSync = final_sync_zs;

    my_data.finalColorImage = FF_NULL_HANDLE;
    //my_data.finalColorInitialSync = initial_sync;
    //my_data.finalColorFinalSync = final_sync;

    return IsInit() ? FrameFlowVK_ProvidePerFrameData(ffInstance_, &my_data) : FFResult::FF_ERROR_NOT_INITIALIZED;
}

FFResult FFWrapperVK::OnDrawFrame(void* tex_handle) noexcept
{
    FFLOGI("Sdk::Unity::FFWrapperVK::OnDrawFrame - Start");
    if (!IsInit()) {
        return FFResult::FF_ERROR_NOT_INITIALIZED;
    }

    unity_graphics_->EnsureOutsideRenderPass();

    FFVKImageSync initial_sync{};
    initial_sync.accessMask = VK_ACCESS_SHADER_WRITE_BIT;
    initial_sync.layout = VK_IMAGE_LAYOUT_COLOR_ATTACHMENT_OPTIMAL;
    initial_sync.stages = VK_PIPELINE_STAGE_FRAGMENT_SHADER_BIT;

    FFVKImageSync final_sync{};
    final_sync.accessMask = VK_ACCESS_SHADER_READ_BIT;
    final_sync.layout = VK_IMAGE_LAYOUT_SHADER_READ_ONLY_OPTIMAL;
    final_sync.stages = VK_PIPELINE_STAGE_FRAGMENT_SHADER_BIT;

    auto register_result = RegisterImage(tex_handle, false, initial_sync, g_unity_predict);
    if (!register_result) {
        FFLOGE("FFWrapperVK::OnProvidePerFrameData. Cant register predict texture.");
        return FF_ERROR_INTERNAL_ERROR;
    }

    UnityVulkanRecordingState recordingState{};
    bool recording_state = unity_graphics_->CommandRecordingState(&recordingState, kUnityVulkanGraphicsQueueAccess_DontCare);
    if (!recording_state) {
        FFLOGE("Sdk::Unity::FFWrapperVK::OnDrawFrame - Cant set recording state");
        return FF_ERROR_INTERNAL_ERROR;
    }

    FFVKDrawFrameInfo info{};
    info.sType = FF_STRUCTURE_TYPE_DRAW_FRAME_INFO;
    info.vkCommandBuffer = recordingState.commandBuffer;
    info.frameIdx = (++m_frame_idx) % FRAMES_IN_FLIGHT;
    info.targetImage = g_unity_predict;
    info.targetInitialSync = initial_sync;
    info.targetFinalSync = final_sync;
    info.doDraw = FF_TRUE;

    FFResult const result = FrameFlowVK_DrawFrame(ffInstance_, &info);

    if (result != FF_SUCCESS) {
        FFLOGW("Sdk::Unity::FFWrapperVK::OnDrawFrame - FrameFlow_DrawFrame returned %d", static_cast<int>(result));
    }

    return result;
}

bool FFWrapperVK::IsInit() const noexcept
{
    return ffInstance_ != FF_NULL_HANDLE;
}

bool FFWrapperVK::LoadFrameFlow() noexcept
{
    FFLOGI("FFWrapperVK::LoadFrameFlow - start");
    if (frameflowLib_.IsLoaded()) {
        return true;
    }
    return frameflowLib_.Load("libframeflow.so", FF::Platform::LibraryLoadPolicy::LOCAL);
}

void FFAPI_PTR FFWrapperVK::SDKErrorCallback(char const *message, void * /*userData*/)
{
    FFLOGE("Sdk::Unity::FFWrapperVK::SDKErrorCallback - FrameFlow SDK: %s", message);
}

} // namespace FFSdkUnity::Vulkan
